package com.caochenlei.graph;

import java.util.LinkedList;
import java.util.PriorityQueue;
import java.util.Queue;

public class KruskalMST {
    private Queue<Edge> mst;                    //用于存储最小生成树的各个边
    private Queue<Edge> pq;                     //用于存储对图各边按权重排序
    private int[] eleAndGroup;                  //下标代表顶点，值代表哪个组

    public KruskalMST(EdgeWeightedGraph G) {
        this.mst = new LinkedList<>();          //初始化mst队列
        this.pq = new PriorityQueue<>();        //初始化pq队列
        this.eleAndGroup = new int[G.size()];   //初始化eleInGroup
        //使用优先队列pq按照权重值排序所有的边
        for (Edge e : G.edges()) {
            pq.add(e);
        }
        //将G.size()个顶点分为G.size()个组
        for (int i = 0; i < eleAndGroup.length; i++) {
            eleAndGroup[i] = i;
        }
        //从优先队列中不停弹出最小边，然后判断
        while (!pq.isEmpty() && (mst.size() < G.size() - 1)) {
            Edge e = pq.poll();                 //取出权重最小的那条边
            int v = e.either();                 //获取权重最小的那条边的其中一个点
            int w = e.other(v);                 //获取权重最小的那条边的另外一个点
            int vGroup = eleAndGroup[v];        //获取顶点v所在的分组，下标v代表顶点v，值代表顶点v所在的组号
            int wGroup = eleAndGroup[w];        //获取顶点w所在的分组，下标w代表顶点w，值代表顶点w所在的组号
            if (vGroup == wGroup) {             //如果顶点v和顶点w在同一分组中
                continue;                       //如果在同一个组中，说明是无效边，则应该直接跳过
            }
            //如果顶点v和顶点w不在同一分组中，说明这是一条有效的权重最小的边
            //为了方便接下来对无效边的判断，我们应该把顶点v和顶点w放到同一组
            for (int i = 0; i < eleAndGroup.length; i++) {
                if (eleAndGroup[i] == vGroup) {
                    eleAndGroup[i] = wGroup;
                }
            }
            mst.add(e);                         //将权重最小的那条边加入到最小生成树中
        }
    }

    //获取最小生成树的各边
    public Queue<Edge> edges() {
        return mst;
    }

    //获取最小生成树的权重
    public double weight() {
        double s = 0.0D;
        if (mst != null) {
            for (Edge e : mst) {
                s += e.weight();
            }
        }
        return s;
    }
}
